#!/usr/bin/env python3
import os
import subprocess
import re
import sys
import time  # 新增：用于计时

# One-hot code mappings
P_ORIG_CODES   = ['100', '010', '001']  # partition_type for orig: cyclic, block, complete
P_SOL_CODES    = ['10', '01']           # partition_type for sol: block, complete
P_FILTER_CODES = ['10', '01']           # partition_type for filter: block, complete
PF_ORIG_CODES  = ['100', '010', '001']  # partition_factor for orig: 8, 4, 2 
UNROLL_CODES   = ['100', '010', '001']  # outer/inner 循环展开路数: 8, 4, 2

# Mapping functions
def map_partition_orig(code):
    return {'100': 'cyclic', '010': 'block', '001': 'complete'}[code]

def map_partition_sol_filter(code):
    return {'10': 'block', '01': 'complete'}[code]

def map_factor(code):
    return {'100': '8', '010': '4', '001': '2'}[code]


def modify_config(config_path, p_orig, p_sol, p_filter, pf_orig, outer_unroll, inner_unroll):
    """Modify the Aladdin config file in place according to provided codes."""
    with open(config_path, 'r') as f:
        lines = f.readlines()

    updated = []
    for line in lines:
        parts = line.strip().split(',')
        if parts[0] == 'partition':
            target = parts[2]
            if target == 'orig':
                parts[1] = map_partition_orig(p_orig)
                parts[5] = map_factor(pf_orig)
            elif target == 'sol':
                parts[1] = map_partition_sol_filter(p_sol)
            elif target == 'filter':
                parts[1] = map_partition_sol_filter(p_filter)
        elif parts[0] == 'unrolling':
            loop = parts[2]
            if loop == 'outer':
                parts[3] = map_factor(outer_unroll)
            elif loop == 'inner':
                parts[3] = map_factor(inner_unroll)
        updated.append(','.join(parts) + '\n')

    with open(config_path, 'w') as f:
        f.writelines(updated)


def run_aladdin(trace_path, config_path, summary_path):
    """Run Aladdin and capture stdout to summary_path."""
    aladdin_home = os.environ.get('ALADDIN_HOME')
    if not aladdin_home:
        print('Error: ALADDIN_HOME not set', file=sys.stderr)
        return False
    exe = os.path.join(aladdin_home, 'common', 'aladdin')
    if not os.path.isfile(exe):
        print(f'Error: aladdin executable not found at {exe}', file=sys.stderr)
        return False

    # Execute Aladdin, capture stdout
    result = subprocess.run(
        [exe, 'stencil', trace_path, config_path],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True
    )
    if result.returncode != 0:
        print(f'Error: Aladdin failed: {result.stderr}', file=sys.stderr)
    # Write stdout to summary file
    with open(summary_path, 'w') as f:
        f.write(result.stdout)
    return True


def parse_summary(summary_path):
    """Extract cycle, avg power, total area from summary file."""
    try:
        text = open(summary_path).read()
    except FileNotFoundError:
        return '', '', ''

    cycle_match = re.search(r"Cycle[s]?[\s:]+([0-9,.]+)", text, re.IGNORECASE)
    power_match = re.search(r"avg\s*power[\s:]+([0-9.eE+-]+)", text, re.IGNORECASE)
    area_match  = re.search(r"total\s*area[\s:]+([0-9.eE+-]+)", text, re.IGNORECASE)

    cycle = cycle_match.group(1) if cycle_match else 'N/A'
    power = power_match.group(1) if power_match else 'N/A'
    area  = area_match.group(1)  if area_match  else 'N/A'
    return cycle, power, area


def main():
    # 计时开始
    start_time = time.time()

    # Configuration
    os.chdir("/workspace/gem5-aladdin/src/aladdin/SHOC/stencil/example/tmp/block")  # 根据实际路径修改

    config_path  = './config_example'
    trace_path   = './dynamic_trace.gz'
    summary_file = './stencil_summary'
    result_file  = './res.txt'

    # Initialize result file
    with open(result_file, 'w') as rf:
        rf.write('# p_orig,p_sol,p_filter,pf_orig,outer,inner, Cycle, avg power, total area\n')

    total = (len(P_ORIG_CODES) * len(P_SOL_CODES) * len(P_FILTER_CODES) *
             len(PF_ORIG_CODES) * len(UNROLL_CODES) * len(UNROLL_CODES))
    count = 0

    for p_orig in P_ORIG_CODES:
        for p_sol in P_SOL_CODES:
            for p_filter in P_FILTER_CODES:
                for pf_orig in PF_ORIG_CODES:
                    for outer in UNROLL_CODES:
                        for inner in UNROLL_CODES:
                            count += 1
                            param_str = f"{p_orig},{p_sol},{p_filter},{pf_orig},{outer},{inner}"
                            print(f"[{count}/{total}] Running: {param_str}")

                            modify_config(config_path, p_orig, p_sol, p_filter, pf_orig, outer, inner)
                            success = run_aladdin(trace_path, config_path, summary_file)
                            if not success:
                                cycle = power = area = 'ERROR'
                            else:
                                cycle, power, area = parse_summary(summary_file)

                            with open(result_file, 'a') as rf:
                                rf.write(f"{param_str}, {cycle}, {power}, {area}\n")

    # 计算并打印总用时
    elapsed = time.time() - start_time
    print(f"Done! Results saved to {result_file}")
    print(f"Total elapsed time: {elapsed:.2f} seconds")

if __name__ == '__main__':
    main()
